[RadeonFlow] flydsl mxfp4 a4w4 MoE: bring gemm1/gemm2 to parity with the HIP backend - #3828
Merged
coderfeli merged 3 commits intoJun 22, 2026
Merged
Conversation
Pass bare data_ptr() device addresses (fx.Int64) for the global buffer args instead of full memref descriptors. The kernels only need base pointers (contiguity + compile-time sizes), so the addresses pack contiguously into kernargs -> coalesced s_load prologue. ~7% faster at decode (M=4/8), converging to parity by M=64. cos preserved.
- backend-managed K-loop waitcnt: drop the hand-tuned inline-asm vmcnt that the LLVM waitcnt pass double-inserted (the large-M scheduling bubble); let rocdl.barrier() own the load->ds_read wait. - epilog output-address strength reduction: hoist the loop-invariant row*N_OUT i64 multiply out of the per-element store loop; per-element offsets become compile-time constants folded into the store address. - unsigned index division (divui/remui) for the non-negative grid/tile/count index math, dropping signed-division sign-correction SALU. M=4096 1.11 -> 1.00 (parity with HIP), M=16384 1.10 -> 1.03; faster/parity through M=2048 unchanged. cos preserved.
Use rocdl.exp2 (v_exp_f32) in silu_mul instead of the software math.exp2 expansion (matches HIP __expf). Removes ~31% of non-MFMA VALU (the v_ldexp + range-clamp v_cmp/v_cndmask the polynomial emitted); large-M ~2-6% faster (M=4096 1.07 -> 1.05). cos preserved.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Brings the flydsl-ported mxfp4 (a4w4) MoE gemm1/gemm2 kernels on gfx950 up to the
HIP backend's performance: closes the large-M (prefill) gap to parity while keeping
the decode/mid-batch lead. Numerically identical (cos ≥ 0.999991 vs HIP, =1 at large M).
Changes
ABI (both kernels + their wrappers)
data_ptr()(int64) instead of afull memref/
fx.Tensordescriptor. Kernels assume contiguity and derive sizes fromcompile-time constants, so descriptor shape/stride fields are dead weight; bare
addresses let the kernarg
s_loads coalesce. ~7% faster at decode, parity by M≈64.gemm2 (
aiter/ops/flydsl/kernels/mxfp4_gemm2.py)inline-asm
s_waitcnt vmcnt(N)and let the LLVM waitcnt pass own it. The inlineasm was opaque to that pass, which double-inserted its own waitcnt — the large-M
scheduling bubble.
row*N_OUTi64 multiply out of the per-element store loop; per-element offsets become
compile-time constants folded into the store address.
(drops signed-division sign-correction SALU).
gemm1 (
aiter/ops/flydsl/kernels/mxfp4_gemm1.py)exp2insilu_mul(rocdl.exp2→v_exp_f32) instead of thesoftware
math.exp2expansion (matches HIP's__expf).Measurement
Per-gemm device time via a cudagraph microbench: capture N back-to-back gemm
launches into one hipgraph, replay, and divide the event-timed elapsed by N
(isolates the gemm from the surrounding pipeline). Exclusive GPU, several
interleaved trials, reported as the in-run flydsl/hip ratio.
Results
Kimi-K2.5 TP=4 shape (NE=385, D_HIDDEN=7168, D_INTER=512, TOPK=9), gfx950,
per-gemm device time, flydsl/hip ratio (lower = faster than HIP):
Notes
preserved at all M.
TODO
~4% (MFMA/K-loop bound, K = D_HIDDEN = 7168 dominates), gemm2 ~3% at M = 16384.